import re
import networkx as nx
from typing import List
import community as community_louvain
from pydantic import BaseModel, Field
from src.common.logger import getLogger
from src.agentic.config.SpacyModel import SpacyModel
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

logger = getLogger()

class Entity(BaseModel):
    name: str = Field(description = "实体名")
    type: str = Field(description = "实体类型")
    desc: str = Field(description = "描述")

class Relation(BaseModel):
    source: str = Field(description="源实体")
    type: str = Field(description="关系类型")
    target: str = Field(description="目标实体")

class EntityRelation(BaseModel):
    entities: List[Entity] = Field(description = "实体")
    relations: List[Relation] = Field(description = "关系")

class CommunityGraph:

    def __init__(self, llm_model, embed_model, llm_or_spacy: bool = False):
        self.llm_model = llm_model
        self.embed_model = embed_model
        self.graph_nx = nx.Graph()
        self.llm_or_spacy = llm_or_spacy
        self.spacy_model = SpacyModel()

    def build_knowledge_graph(self, documents):
        logger.info(f"CommunityGraph build_knowledge_graph documents len: {len(documents)}")
        for document in documents:
            entities, relations = self.extract_entities_relations(document)
            for entity in entities:
                self.graph_nx.add_node(entity["name"], name = entity["name"], type = entity["type"], desc = entity["desc"])

            for relation in relations:
                self.graph_nx.add_edge(relation["source"], relation["target"], relation = relation["type"])
        logger.info("CommunityGraph build_knowledge_graph finish")

    def extract_entities_relations(self, document):
        logger.info(f"CommunityGraph extract_entities_relations document len: {len(document)}")
        if self.llm_or_spacy:
            return self.extract_entity_relation(document)
        else:
            return self.spacy_model.extract_entity_relations(document)

    def extract_entity_relation(self, document):
        logger.info(f"CommunityGraph extract_entity_relation document len: {len(document)}")
        template = """
            你是一个专业的信息抽取系统。请从以下文本中精确识别并提取所有有意义的命名实体。

            要求：
            1. 只提取明确提及的实体，不要推测或生成不存在的内容。
            2. 实体类型包括但不限于：人物（Person）、组织（Organization）、地点（Location）、日期/时间（Date）、技术术语（Technology）、产品（Product）、事件（Event）。
            3. 对每个实体，提供实体在文本中的原始形式（保持大小写和标点）
            4. 不要包含任何解释、前缀、后缀或 Markdown。

            文本： {context}
        """
        prompt = ChatPromptTemplate.from_template(template)
        extract_chain = prompt | self.llm_model.with_structured_output(EntityRelation)
        extract_result = extract_chain.invoke({ "context": document })
        entities = extract_result.entities
        logger.info(f"CommunityGraph extract_entity_relation entities len: {len(entities)}")
        relations = extract_result.relations
        logger.info(f"CommunityGraph extract_entity_relation relations len: {len(relations)}")
        return entities, relations

    def detect_community_graph(self):
        logger.info("CommunityGraph detect_community_graph start")
        communities = community_louvain.best_partition(self.graph_nx)
        logger.info(f"CommunityGraph detect_community_graph communities len: {len(communities)}")

        community_nodes = {}
        for node, community_id in communities.items():
            if community_id not in community_nodes:
                community_nodes[community_id] = []
            community_nodes[community_id].append(node)
        return community_nodes

    def generate_community_summary(self, communities):
        logger.info("CommunityGraph generate_community_summary start")
        community_summaries = {}
        for community_id, nodes in communities.items():
            logger.info(f"CommunityGraph generate_community_summary community_id: {community_id}, nodes: {len(nodes)}")
            subGraph = self.graph_nx.subgraph(nodes)

            community_info = f"社区包含{len(nodes)}个实体：{', '.join(nodes)}\n"
            community_info += f"关系包括：\n"

            for edge in subGraph.edges(data = True):
                source, target, attrs = edge
                relation = attrs.get("relation", "相关")
                community_info += f"- {source} {relation} {target}\n"

            community_summaries[community_id] = self.summary_community_info(community_info)
            logger.info(f"CommunityGraph generate_community_summary community_summaries keys len: {len(community_summaries.keys())}")
            return community_summaries

    def summary_community_info(self, community_info):
        logger.info(f"CommunityGraph summary_community_info community_info len: {len(community_info)}")
        template = """
            请为以下知识社区生成简洁的摘要：

            {community_info}

            摘要应该：
            1. 概括社区的主要内容
            2. 突出关键实体和关系
            3. 保持简洁易懂
        """
        prompt = ChatPromptTemplate.from_template(template)
        summary_chain = prompt | self.llm_model | StrOutputParser()
        response = summary_chain.invoke({ "community_info": community_info })
        logger.info(f"CommunityGraph summary_community_info response len: {len(response)}")
        return response

    def find_relevent_community(self, query, community_summaries):
        logger.info("CommunityGraph find_relevent_community start")
        query_words = set(re.findall(r'\w+', query.lower()))
        logger.info(f"CommunityGraph find_relevent_community query_words: {query_words}")

        relevent_communities = []
        for community_id, summary in community_summaries.items():
            summary_words = set(re.findall(r'\w+', summary.lower()))
            logger.info(f"CommunityGraph find_relevent_community summary_words len: {len(summary_words)}")
            overlap = len(query_words & summary_words)
            logger.info(f"CommunityGraph find_relevent_community community_id: {community_id}, overlap: {overlap}")
            if overlap > 0:
                relevent_communities.append((community_id, overlap))
        logger.info(f"CommunityGraph find_relevent_community relevent_communities len: {len(relevent_communities)}")

        relevent_communities.sort(key = lambda x: x[1], reverse = True)
        logger.info("CommunityGraph find_relevent_community finish")
        return [id for id, _ in relevent_communities[:3]]

    def invoke(self, query, documents):
        logger.info("CommunityGraph invoke start")
        self.build_knowledge_graph(documents)

        community_nodes = self.detect_community_graph()
        logger.info(f"CommunityGraph invoke community_nodes keys len: {len(community_nodes.keys())}")

        community_summaries = self.generate_community_summary(community_nodes)
        logger.info(f"CommunityGraph invoke community_summaries keys len: {len(community_summaries.keys())}")

        relevent_communities = self.find_relevent_community(query, community_summaries)
        logger.info(f"CommunityGraph invoke relevent_communities len: {len(relevent_communities)}")

        community_graphs = []
        for community_id in relevent_communities:
            community_graphs.append(community_summaries[community_id])
        logger.info(f"CommunityGraph invoke community_graphs len: {len(community_graphs)}")
        return community_graphs
